pub mod prelude;

use std::{collections::HashMap, sync::Arc};
use crate::{protocol::McpTool, transport::stdio::McpStdioTransport, McpClientInitializeResult, McpError, McpImplementation, McpPromptsCapability, McpResourcesCapability, McpResult, McpServerCapabilities, McpServerInitializeResult, McpToolCallResult, McpToolCallResultContent, McpToolsCapability};
use crate::server::{McpServer, McpServerHandler};
use anyhow::Result;
use async_trait::async_trait;
use log::debug;

#[async_trait]
pub trait Tool: Send + Sync {
    fn name(&self) -> String;
    fn intto_mcptool(&self) -> McpTool;
    async fn execute(&self, args: serde_json::Value) -> Result<Option<String>>;
}

pub enum Transport {
    Stdio,
}

pub struct FastMcp {
    server: McpServer,
}

struct FastMcpHandler {
    tools: HashMap<String, Box<dyn Tool>>
}

impl FastMcpHandler {
    fn new() -> Self {
        Self { tools: HashMap::new() }
    }
}

#[async_trait]
impl McpServerHandler for FastMcpHandler {
    async fn initialize(&self, result: McpClientInitializeResult) -> McpResult<McpServerInitializeResult> 
    {
        debug!("Client connected: {} v{}", result.client_info.name, result.protocol_version);
        
        Ok(McpServerInitializeResult {
            protocol_version: crate::protocol::LATEST_PROTOCOL_VERSION.to_string(),
            capabilities: McpServerCapabilities {
                experimental: Some(HashMap::new()),
                logging: None,
                prompts: Some(McpPromptsCapability {
                    list_changed: Some(false),
                }),
                resources: Some(McpResourcesCapability {
                    subscribe: Some(false),
                    list_changed: Some(false),
                }),
                tools: Some(McpToolsCapability {
                    list_changed: Some(false)
                })  
            },
            server_info: McpImplementation {
                name: "Hello".to_string(),
                version: "1.0.1".to_string()
            },
            instructions: None
        })
    }

    async fn tools_list(&self) -> McpResult<Vec<McpTool>> {
        Ok(self.tools.iter().map(|(_, tool)| tool.intto_mcptool()).collect())
    }

    async fn execute_tool(
        &self,
        tool_name: &str,
        arguments: Option<serde_json::Value>,
    ) -> McpResult<McpToolCallResult> {
        match self.tools.get(tool_name) {
            Some(tool) => {
                let arguments = arguments.unwrap_or(serde_json::json!({}));
                match tool.execute(arguments).await {
                    Ok(result) => {
                        let content = match result {
                            Some(ret_value) => vec![ McpToolCallResultContent::Text { text: ret_value } ],
                            None => vec! []
                        };
                        Ok(McpToolCallResult{
                            content,
                            is_error: Some(false)
                        })
                    }
                    Err(err) => Err(McpError::Other(err.to_string()))
                }

            }
            None => Err(McpError::Other(format!("No exit tool '{}'", tool_name)))
        }
    }

    async fn shutdown(&self) -> McpResult<()> {
        debug!("Server shutting down");
        Ok(())
    }
}

impl FastMcp {
    pub fn new(transport: Transport, tools: Vec<Box<dyn Tool>>) -> Self {
        let mut fast_handler = FastMcpHandler::new();
        for tool in tools {
            fast_handler.tools.insert(tool.name(), tool);
        }

        match transport {
            Transport::Stdio => {
                let (cmp_transport, _) = McpStdioTransport::new();
                Self {
                    server: McpServer::new(Arc::new(cmp_transport), Arc::new(fast_handler))
                } 
            }
        }
    }

    pub async fn run(self) -> McpResult<()> {
        let handle = tokio::spawn(async move {
            if let Err(e) = self.server.run().await {
                eprintln!("Error in connection: {}", e);
            }
        });
        handle.await.unwrap();
        Ok(())
    }
}


